import json
import sys
from ipwhois import IPWhois
import requests


class IPSearch:
    def __init__(self):
        pass

    def search_ip_whois(self, ip):
        """
        使用ipwhois库查询IP归属信息
        """
        try:
            obj = IPWhois(ip)
            results = obj.lookup_rdap(depth=1)
            return results
        except Exception as e:
            return {"error": f"查询过程中发生错误: {str(e)}"}

    def format_whois_result(self, result):
        """
        格式化WHOIS查询结果
        """
        if "error" in result:
            return result["error"]

        formatted = {
            "IP地址": result.get('query', 'N/A'),
            "国家": result.get('asn_country_code', 'N/A'),
            "ASN": result.get('asn', 'N/A'),
            "ASN注册机构": result.get('asn_registry', 'N/A'),
            "ASN描述": result.get('asn_description', 'N/A'),
            "网络名称": result.get('network', {}).get('name', 'N/A') if result.get('network') else 'N/A',
            "网络起始地址": result.get('network', {}).get('start_address', 'N/A') if result.get('network') else 'N/A',
            "网络结束地址": result.get('network', {}).get('end_address', 'N/A') if result.get('network') else 'N/A',
            "网络描述": result.get('network', {}).get('description', 'N/A') if result.get('network') else 'N/A',
        }
        return formatted

    def search_ip_api(self, ip):
        """
        使用免费API查询IP归属信息
        """
        try:
            # 使用ip-api.com的免费API
            url = f"http://ip-api.com/json/{ip}"
            response = requests.get(url, timeout=5)
            if response.status_code == 200:
                data = response.json()
                return {
                    "IP地址": data.get("query", "N/A"),
                    "国家": data.get("country", "N/A"),
                    "国家代码": data.get("countryCode", "N/A"),
                    "省份": data.get("regionName", "N/A"),
                    "城市": data.get("city", "N/A"),
                    "ISP": data.get("isp", "N/A"),
                    "组织": data.get("org", "N/A"),
                    "AS": data.get("as", "N/A"),
                    "纬度": data.get("lat", "N/A"),
                    "经度": data.get("lon", "N/A"),
                    "时区": data.get("timezone", "N/A")
                }
            else:
                return {"error": f"API请求失败，状态码: {response.status_code}"}
        except Exception as e:
            return {"error": f"API查询过程中发生错误: {str(e)}"}

    def search_ip(self, ip, method="whois"):
        """
        查询IP地址归属信息的主函数
        :param ip: 要查询的IP地址
        :param method: 查询方法 ("whois" 或 "api")
        :return: 查询结果
        """
        if method == "whois":
            result = self.search_ip_whois(ip)
            return self.format_whois_result(result)
        elif method == "api":
            return self.search_ip_api(ip)
        else:
            return {"error": "不支持的查询方法，请使用 'whois' 或 'api'"}


def main():
    # 创建IPSearch实例
    ip_search = IPSearch()
    
    # 如果命令行提供了IP地址，则使用它
    if len(sys.argv) > 1:
        ip = sys.argv[1]
        print(f"正在查询IP地址: {ip}")
        
        # 使用WHOIS方法查询
        print("\n=== 使用WHOIS方法查询 ===")
        result = ip_search.search_ip(ip, method="whois")
        print(json.dumps(result, ensure_ascii=False, indent=2))
        
        # 使用API方法查询
        print("\n=== 使用API方法查询 ===")
        result = ip_search.search_ip(ip, method="api")
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        # 交互式查询
        print("=== IP地址归属查询工具 ===")
        while True:
            ip = input("\n请输入要查询的IP地址 (输入 'quit' 退出): ").strip()
            if ip.lower() == 'quit':
                break
            
            if not ip:
                print("请输入有效的IP地址")
                continue
                
            print(f"\n正在查询IP地址: {ip}")
            
            # 提供选择查询方法
            method = input("请选择查询方法 (1: WHOIS, 2: API, 3: 两种都查询) [默认为API]: ").strip()
            if method == "1":
                result = ip_search.search_ip(ip, method="whois")
                print("\n=== 查询结果 (WHOIS) ===")
                print(json.dumps(result, ensure_ascii=False, indent=2))
            elif method == "2":
                result = ip_search.search_ip(ip, method="api")
                print("\n=== 查询结果 (API) ===")
                print(json.dumps(result, ensure_ascii=False, indent=2))
            elif method == "3":
                print("\n=== 查询结果 (WHOIS) ===")
                result = ip_search.search_ip(ip, method="whois")
                print(json.dumps(result, ensure_ascii=False, indent=2))
                
                print("\n=== 查询结果 (API) ===")
                result = ip_search.search_ip(ip, method="api")
                print(json.dumps(result, ensure_ascii=False, indent=2))
            else:
                result = ip_search.search_ip(ip, method="api")
                print("\n=== 查询结果 (API) ===")
                print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()